import sys
sys.path.append("..")
from abc import abstractmethod
import requests
import time
import os
import whisper

from .utils import gen_random_job_id, base64_to_audio_file

class BaseASR:
    def __init__(self):
        ...

    @abstractmethod
    def audio2txt(self, audio: str) -> str:
        raise NotImplementedError("Please implemented audio2txt function.")

    def __call__(self, audio: str) -> str:
        text = self.audio2txt(audio).strip()
        return text


class WhisperASR(BaseASR):
    def __init__(self):
        super(WhisperASR, self).__init__()
        self.whisper_model = whisper.load_model("turbo", download_root="model")
        self.tmp_save_dir = "./tmp"
        if not os.path.exists(self.tmp_save_dir):
            os.makedirs(self.tmp_save_dir, exist_ok=True)

    def audio2txt(self, audio):
        tmp_imgpath = os.path.join(self.tmp_save_dir, f"{gen_random_job_id()}.mp3")
        if audio.startswith(('http://', 'https://')): # 处理网络图片
            response = requests.get(audio)
            with open(tmp_imgpath, 'wb') as f:
                f.write(response.content)
        elif len(audio) > 100:  # 此时为base64编码
            tmp_imgpath = base64_to_audio_file(audio, tmp_imgpath)
        # 这两个条件都不满足那么就是普通本地路径
        result = self.whisper_model.transcribe(tmp_imgpath)["text"]
        os.remove(tmp_imgpath)
        return result
    
if __name__ == "__main__":
    asr_model = WhisperASR()
    audio = "test.mp3"
    res = asr_model(audio)
    print(res)
